import os
import torch
import torchvision
from torch import nn
# from d2l import torch as d2l


data_dir = os.path.join('data', '')
assert os.path.exists(data_dir), f'data_dir[{data_dir}]不存在?'

train_imgs = torchvision.datasets.ImageFolder(os.path.join(data_dir, 'train'))
test_imgs = torchvision.datasets.ImageFolder(os.path.join(data_dir, 'test'))

pretrained_net = torchvision.models.resnet18(pretrained=True)
pretrained_net.fc

finetune_net = torchvision.models.resnet18(pretrained=True)
finetune_net.fc = nn.Linear(finetune_net.fc.in_features, 2)
nn.init.xavier_uniform_(finetune_net.fc.weight)


scratch_net = torchvision.models.resnet18()
scratch_net


